import chromadb.utils.embedding_functions as embedding_functions
import chromadb
import json

ollama_ef = embedding_functions.OpenAIEmbeddingFunction(
                api_key="sk-lOnR4erWFrXwpBqW24E293122a0347AdA983D7E99e0512B1",
                api_base="http://60.204.226.75:32100/v1",
                model_name="mxbai-embed-large"
            )
documents=[]
# documents = [
#         {'报警历史表': 'alarm.alarm_history', '字段信息': 'id, lamp_post_id, device_id, alarm_rule_id, content, first_alarm_time, last_alarm_time, insert_time, update_time'},
#         {'报警规则表': 'alarm.alarm_rule', '字段信息': 'id, alarm_type, rule_name, alarm_level, trigger_expression, time_span, detection_cycle, alarm_content, is_enabled, organ_ids, insert_time, update_time, comment, alarm_event_type, parameter_lang'},
#         {'报警规则分组': 'alarm.alarm_rule_trigger_group', '字段信息': 'id, alarm_rule_id, sort'},
#         {'报警规则明细': 'alarm.alarm_rule_trigger_group_item', '字段信息': 'id, rule_trigger_group_id, attribute_field, operator, value, insert_time, update_time, sort'},
#         {'报警类型与行业类型对应关系': 'camera.alarm_type_industry_relation', '字段信息': 'id, aid, code, name'},
#         {'算法报警记录表': 'camera.algorithm_alarm', '字段信息': 'id, aid, cid, cid_url, status, time_stamp, srcpic_name, srcpic_data, pic_name, pic_data, data, update_time, insert_time, process_status, alarm_type, beike_sound_id'},
#         {'报警配置': 'camera.algorithm_alarm_config', '字段信息': 'id, alarm_type, status, start_time, end_time, remarks, update_time, insert_time, alarm_code'},
#         {'算法报警联动表': 'camera.algorithm_alarm_linkage', '字段信息': 'id, cid, aid, linkage_device_type, linkage_device_id, content, update_time, insert_time, linkage_alarm_device_type, linkage_type'},
#         {'算法报警联动动作明细': 'camera.algorithm_alarm_linkage_item', '字段信息': 'id, alarm_linkage_id, linkage_action_type, update_time, insert_time, parameter, sort'},
#         {'报警处理表': 'camera.algorithm_alarm_process', '字段信息': 'id, algorithm_alarm_id, handler, process_method, process_pic, content, process_status, update_time, insert_time'},
#         {'算法报警记录表': 'camera.algorithm_camera', '字段信息': 'id, aids, cid, update_time, insert_time, location, code, ip_port, platform_type, info'},
#         {'单灯控制器方案结果': 'light.alonelamp_control_program_record', '字段信息': 'id, imei, info, insert_time, update_time, thiot_device_id'},
#         {'灯源用途管理': 'device.b_alonelamp_effect', '字段信息': 'lamp_effect_id, lamp_effect_name, tenant_id, update_time, insert_time'},
#         {'行政区域表': 'device.b_canton', '字段信息': 'canton_id, canton_code, canton_name, x_tude, y_tude, parent_id, remark, wgs84x_tude, wgs84y_tude, tenant_id'},
#         {'灯杆类型表': 'device.b_lamppost_type', '字段信息': 'lamp_post_type_id, lamp_post_type_name, alone_lamp_count, tenant_id, insert_time, update_time'},
#         {'道路信息表': 'device.b_street', '字段信息': 'street_id, street_name, name_spell, start_street_id, end_street_id, tenant_id'},
#         {'灯杆路侧': 'device.b_street_side', '字段信息': 'street_side_id, street_side_name, street_side_number, tenant_id, insert_time, update_time'},
#         {'道路和路侧关联表': 'device.b_street_street_side', '字段信息': 'street_side_relation_id, street_id, street_side_id'},
#         {'音柱设备播放任务': 'broadcast.broadcast_cop_device', '字段信息': 'devicecode, state, file, taskid, broadcastcopdeviceid, eventime, tenant_id'},
#         {'音柱播放历史记录': 'broadcast.broadcast_his_device', '字段信息': 'broadcasthisdeviceid, taskid, taskname, eventime, playdurations, status, publisher, receiveequipment, triggerinfo, commandinfo, tasklevel, tenant_id'},
#         {'摄像头预置位信息表': 'camera.camera_preset', '字段信息': 'id, code, preset_num, preset_name, update_time, insert_time'},
#         {'视频抓图': 'camera.camera_snap', '字段信息': 'id, code, pic_data, update_time, insert_time'},
#         {'厂商信息表': 'device.company', '字段信息': 'id, name, tenant_id, insert_time, update_time, comment'},
#         {'控制器单元过滤表': 'light.control_unit_filter', '字段信息': 'filter_id, street_filter, street_side_filter, lamp_effect_filter, lamppost_number_filter, insert_time, update_time'},
#         {'公共广播信息表': 'device.device_broadcast', '字段信息': 'id, lamp_post_id, name, code, broadcast_out_volume, device_model_id, company_id, online_state, register_state, insert_time, update_time, comment, organ_id'},
#         {'视频设备管理': 'device.device_camera', '字段信息': 'id, lamp_post_id, type, name, code, ip, port, username, password, web_url, connect_type, device_model_id, company_id, register_state, insert_time, update_time, comment, organ_id, mobile_url, online_state, platform_code, ptz, video_server'},
#         {'视频服务器管理': 'device.device_camera_server', '字段信息': 'id, name, ip, port, username, password, connect_type, tenant_id, insert_time, update_time, comment, application_name, status'},
#         {'环境监测设备管理': 'device.device_env', '字段信息': 'id, lamp_post_id, name, code, imei, device_model_id, company_id, thiot_device_id, online_state, register_state, insert_time, update_time, comment, organ_id, comm_protocol_name, model_identification, thiot_connect_type, platform_type'},
#         {'智慧锁与灯杆关系表': 'device.device_lamp_post_lock_relation', '字段信息': 'id, lamp_post_id, insert_time, update_time, lock_id'},
#         {'井盖信息': 'device.device_manhole_cover', '字段信息': 'id, cover_code, pic_data, camera_code, insert_time, update_time, comment'},
#         {'产品类型信息表': 'device.device_model', '字段信息': 'id, name, type, code, company_id, tenant_id, insert_time, update_time, comment'},
#         {'信息发布设备信息表': 'device.device_screen', '字段信息': 'id, lamp_post_id, name, code, screen_size, ip, port, device_model_id, company_id, insert_time, update_time, comment, organ_id, screen_type, env_sensor'},
#         {'一键报警管理': 'device.device_talk', '字段信息': 'id, lamp_post_id, name, code, talk_out_volume, device_model_id, company_id, online_state, register_state, insert_time, update_time, comment, organ_id, binding_target_code'},
#         {'无线设备信息表': 'device.device_wireless', '字段信息': 'id, lamp_post_id, name, mac, device_model_id, company_id, online_state, insert_time, update_time, comment, organ_id, ip'},
#         {'灯杆系统字典信息表': 'sysmgr.dic_info', '字段信息': 'id, dic_type_code, name, code, sort, insert_time, update_time, comment'},
#         {'字典类型信息': 'sysmgr.dic_type', '字段信息': 'id, name, code, sort, insert_time, update_time, comment'},
#         {'环境监测上报历史数据表': 'env.env_history_data', '字段信息': 'id, thiot_device_id, tsp, pm_two_dot_five, pm_ten, noise, wind_direction, wind_speed, temperature, humidity, start_time, end_time, insert_time, update_time, pressure'},
#         {'环境监测实时数据表': 'env.env_realtime_data', '字段信息': 'id, thiot_device_id, tsp, pm_two_dot_five, pm_ten, noise, wind_direction, wind_speed, temperature, humidity, start_time, end_time, insert_time, update_time, pressure'},
#         {'灯源管理': 'device.l_alonelamp', '字段信息': 'alone_lamp_id, lamp_post_id, alone_lamp_number, lamp_effect_id, energy_type_id, lamp_type_id, lamp_power, alone_lamp_code, alone_lamp_name, update_time, insert_time, dimming_value'},
#         {'单灯控制器': 'device.l_alonelamp_control', '字段信息': 'alone_lamp_control_id, lamp_post_id, product_number, manufacturer_id, alone_lamp_control_type, alone_lamp_control_model, alone_lamp_protocol_type, imei, comm_protocol_name, model_identification, thiot_connect_type, thiot_device_id, register_state, update_time, insert_time, app_info_id'},
#         {'单灯控制器中间表': 'device.l_alonelamp_control_middle', '字段信息': 'id, alone_lamp_control_id, alone_lamp_id'},
#         {'控制器数据上报记录表': 'light.l_alonelamp_control_property', '字段信息': 'alonelamp_control_property_id, alonelamp_control_id, software_version, hardware_version, protocol_version, heart_beat_cycle, disable_program, imei, imsi, is_scramble, peak_time, curr_thr, vol_range_percent, curr_range_percent, rsrp, ecl, snr, csq, pci, is_network_param, photoctrlmode, photothreshold, photosection, photosmoothing, bluetooth, apn, reselectionswitch, retrycount, timethreshold, barnumberthreshold, device_sn_cur, upgrade_flag, upgrade_ip, upgrade_port, data_center_ip, data_center_port, voltage_frequency, power_threshold, standard_voltage, over_voltage_percent, under_voltage_percent, chip_temperature, abnormal_status, time_plan_status, system_rtc, week'},
#         {'控制器单元信息表': 'light.l_control_unit', '字段信息': 'control_unit_id, onoff_time_id, filter_id, on_off_model, change_light, index, insert_time, update_time'},
#         {'灯源状态': 'light.l_coplog', '字段信息': 'coplog_id, alone_lamp_id, is_online, is_switch_on, warning_info, cop_date, voltage, lamp_isp, power_factor, active_power, reactive_power, plan_open_status, manual_operation_time, lightintensity, leakage_voltage, leakage_current, burning_time, energy, status'},
#         {'巡测历史记录': 'light.l_his_coplog', '字段信息': 'his_coplog_id, alone_lamp_id, lamp_effect_id, cop_date, voltage, lamp_isp, power_factor, active_power, reactive_power, is_online, is_switch_on, warning_info, plan_open_status, lightintensity, leakage_voltage, leakage_current, burning_time, energy, status'},
#         {'灯杆信息表': 'device.l_lamppost', '字段信息': 'lamp_post_id, organ_id, work_area_id, control_cabinet_id, lamp_post_type_id, lamp_post_picture, lamp_post_code, lamp_post_number, alone_lamp_count, add_date, lamp_post_intro, control_count, street_id, street_side_id, x_tude, y_tude, material_type, material_type_id, lamp_height, device_model_id, company_id, category, install_state, name_spell, is_at_crossing, crossing_number, sort_relation, post_code_number, wgs84x_tude, wgs84y_tude, insert_time, update_time, canton_id, smart_light_id'},
#         {'单灯方案列表': 'light.l_onoff_plan', '字段信息': 'onoff_plan_id, onoff_plan_name, remark, insert_time, update_time, organ_id'},
#         {'方案关联道路': 'light.l_onoff_plan_street_group', '字段信息': 'onoff_plan_relation_id, onoff_plan_id, street_group_id'},
#         {'单灯方案时间列表': 'light.l_onoff_time', '字段信息': 'onoff_time_id, onoff_plan_id, start_time, index, insert_time, update_time'},
#         {'方案下发详细列表': 'light.l_program_detail_results', '字段信息': 'program_detail_results_id, on_off_date, start_time, end_time, lamp_post_id, alone_lamp_control_id, alone_lamp_id, light_model, insert_time, update_time, onoff_plan_id'},
#         {'能耗统计': 'light.light_result_electricity_consumption', '字段信息': 'id, time, value, insert_time, update_time, organ_id'},
#         {'亮灯率': 'light.light_result_light_on_rate', '字段信息': 'id, time, value, insert_time, update_time, organ_id'},
#         {'在线率': 'light.light_result_online_rate', '字段信息': 'id, time, value, insert_time, update_time, organ_id'}
#     ]

chroma_client = chromadb.PersistentClient(path="./db")
collection = chroma_client.get_or_create_collection(name="my_demo03_collection", embedding_function=ollama_ef)

for i, d in enumerate(documents):
  embedding = ollama_ef(json.dumps(d,ensure_ascii=False))
  collection.upsert(
    ids=[str(i)],
    embeddings=embedding,
    documents=[json.dumps(d,ensure_ascii=False)]
  )
#query_results = collection.query(query_texts=["灯杆数据"], n_results=3)
query_results = collection.query(query_embeddings=ollama_ef("灯杆信息表"))

print(f"查询结果：{query_results}")